#include <NvInferRuntimeCommon.h>
#include <NvOnnxParser.h>
#include <fstream>
#include <iostream>
#include <opencv2/imgproc.hpp>
#include <sys/stat.h>
#include <glog/logging.h>
#include "yolo_detect_pkg/yolov10.hpp"


using namespace std;
using namespace nvinfer1;

// 定义对
class TRT_Logger : public nvinfer1::ILogger
{
    private:
        nvinfer1::ILogger::Severity _verbosity;
        std::ostream* _ostream;

    public:
        TRT_Logger(Severity verbosity = Severity::kWARNING, std::ostream& ostream = std::cout)
            : _verbosity(verbosity)
            , _ostream(&ostream)
        {
        }
        void log(Severity severity, const char* msg) noexcept override
        {
            if (severity <= _verbosity)
            {
                time_t rawtime = std::time(0);
                char buf[256];
                strftime(&buf[0], 256, "%Y-%m-%d %H:%M:%S", std::gmtime(&rawtime));
                const char* sevstr = (severity == Severity::kINTERNAL_ERROR ? "    BUG" : severity == Severity::kERROR
                            ? "  ERROR"
                            : severity == Severity::kWARNING ? "WARNING" : severity == Severity::kINFO ? "   INFO"
                                                                                                    : "UNKNOWN");
                (*_ostream) << "[" << buf << " " << sevstr << "] " << msg << std::endl;
            }
        }
};

// 判断文件是否形成
static bool ifFileExists(const char *FileName)
{
    struct stat my_stat;
    return (stat(FileName, &my_stat) == 0);
}
// 加载onnx模型
void YOLOv10::loadOnnx(const string strName){
    // 定义输出的日志
    TRT_Logger gLogger;
    // 构建对应的网络
    IBuilder *builder = createInferBuilder(gLogger);
    // 设置对应名处理的Batch——size
    builder->setMaxBatchSize(1);
    const auto explicitBatch = 1U<<static_cast<uint32_t >(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
    // 定义对应的模型
    INetworkDefinition * network = builder->createNetworkV2(explicitBatch);
    // 定义对应onnx模型的解析器
    nvonnxparser::IParser *parser = nvonnxparser::createParser(*network,gLogger);
    // 进行对应的onnx文件解析
    parser->parseFromFile(strName.c_str(),static_cast<int>(ILogger::Severity::kWARNING));
    
    for (int i = 0; i < parser->getNbErrors(); i++)
    {
        cout<<parser->getError(i)->desc()<<endl;
    }
    // 设置模型的配置器
    IBuilderConfig *config = builder->createBuilderConfig();
    // 设置对应的属性
    // 设置对应的空间大小
    config->setMaxWorkspaceSize(1<<20);

    // 来创建一个 ICudaEngine 类型的对象，在构建引擎时，TensorRT会复制权重
    m_CudaEngine=builder->buildEngineWithConfig(*network,*config);
    
    string strTrtName = strName;
    size_t sep_pos = strTrtName.find_last_of(".");
    strTrtName = strTrtName.substr(0,sep_pos)+".engine";
    // 将对应的引擎进行序列化处理
    IHostMemory *gieModelStream = m_CudaEngine->serialize();
    string serialize_str;
    ofstream serialize_output_stream;

    serialize_str.resize(gieModelStream->size());
    memcpy((void *)serialize_str.data(),gieModelStream->data(),gieModelStream->size());
    // 将对应的内容进行输出写出
    serialize_output_stream.open(strTrtName.c_str());
    serialize_output_stream<<serialize_str;
    serialize_output_stream.close();

    m_CudaExecutionContext = m_CudaEngine->createExecutionContext();
    
    parser->destroy();
    network->destroy();
    config->destroy();
    builder->destroy();




}
void YOLOv10::loadTrt(const std::string strName){
    TRT_Logger gLogger;
    // 进行对应的引擎的序列化
    m_CudaRuntime=createInferRuntime(gLogger);
    ifstream fin(strName.c_str());
    std::string cached_engine = "";
    while (fin.peek()!=EOF)
    {
        stringstream buffer;
        buffer<<fin.rdbuf();
        cached_engine.append(buffer.str());
    }
    fin.close();
    // 将对应的模型输入到缓冲区后，将其进行的反序列化
    m_CudaEngine = m_CudaRuntime->deserializeCudaEngine(cached_engine.data(), cached_engine.size(),nullptr);
    // 针对engine进行的上下文内容的执行
    m_CudaExecutionContext = m_CudaEngine->createExecutionContext();
    m_CudaRuntime->destroy();

}

YOLOv10::YOLOv10(){
    std::cout << "yolo init start" << std::endl;
}

void YOLOv10::UnInit()
{

    for(auto &p: m_ArrayDevMemory)
    {      
        cudaFree(p);
        p = nullptr;            
    }        
    for(auto &p: m_ArrayHostMemory)
    {        
        free(p);
        p = nullptr;        
    }        
    cudaStreamDestroy(m_CudaStream);
    

}

YOLOv10::~YOLOv10()
{
    UnInit();   
}


// 进行图像的预处理改变图像大小
cv::Mat YOLOv10::resize_image(cv::Mat srcimg){
    if(srcimg.empty())
    {
        std::cout << "image is empty" << std::endl;
        cv::Mat dst_img = cv::Mat(640, 640, CV_8UC3, Scalar(114, 114, 114));
        return dst_img;
    }
    input_image_height = srcimg.rows;
    input_image_widht= srcimg.cols;
    in_p_height = 640;
    in_p_width = 640;
    newh = in_p_height;
    neww = in_p_width;
    cv::Mat dst_img;
    cv::Mat temp_img;
    if(this->keep_ratio&&input_image_height!=input_image_widht){
        float hw_ratio = (float)input_image_height / (float)input_image_widht;
        if(hw_ratio>1){
            newh = this->in_p_height;
            neww = int(this->in_p_height/hw_ratio);
            cv::resize(srcimg,dst_img,cv::Size(neww,newh),cv::INTER_AREA);
            int paddleft = int((this->in_p_width-neww));
            cv::copyMakeBorder(dst_img,temp_img,0,0,0,paddleft,cv::BORDER_CONSTANT,114);
            
            
        }else{
            newh = (int)this->in_p_height * hw_ratio;
			neww = this->in_p_width;
			cv::resize(srcimg, dst_img, cv::Size(neww, newh), cv::INTER_AREA);
			int paddtop = (int)(this->in_p_height - newh) ;
			cv::copyMakeBorder(dst_img, temp_img, 0, paddtop ,0, 0,cv::BORDER_CONSTANT, 114);
            
        }
    }else{
        cv::resize(srcimg, temp_img, Size(neww, newh), cv::INTER_AREA);
    }
    return temp_img;

}




// 输入图像的归一化处理
void YOLOv10::normalize_(Mat &img){
    int row = img.rows;
    int col = img.cols;
    this->input_image_.resize(row*col*img.channels());
    // bgr
    for(int c = 0; c < 3;c++){
        for(int i= 0;i<row;i++){
            for(int j=0;j<col;j++){
                float pix = img.ptr<uchar>(i)[j*3+2-c];
                this->input_image_[c*row*col+i*col+j] = pix/255.0;
            }
        }
    }
}


void YOLOv10::nms(vector<BoxInfo> &input_boxes){
    if(input_boxes.size() <2){
        return;
    }
    sort(input_boxes.begin(), input_boxes.end(),[](BoxInfo a, BoxInfo b){ return a.score>b.score; });
    vector<float> vArea(input_boxes.size());
    for(int i = 0; i < input_boxes.size();i++){
        vArea[i] = (input_boxes[i].x2-input_boxes[i].x1+1)*(input_boxes[i].y2-input_boxes[i].y1+1);
    }
    vector<bool> isSuppressed(input_boxes.size(),false);
    for(int i = 0; i < input_boxes.size();i++){
        if (isSuppressed[i]) { continue; }
        for (int j = i + 1; i < input_boxes.size(); j++)
        {
            if (isSuppressed[j]) { continue; }
            
			float xx1 = max(input_boxes[i].x1, input_boxes[j].x1);
			float yy1 = max(input_boxes[i].y1, input_boxes[j].y1);
			float xx2 = min(input_boxes[i].x2, input_boxes[j].x2);
			float yy2 = min(input_boxes[i].y2, input_boxes[j].y2);
            float w = max(0.0f, xx2 - xx1 + 1);
			float h = max(0.0f, yy2 - yy1 + 1);
			float inter = w * h;	// 交集
			if(input_boxes[i].label == input_boxes[j].label)
			{
				float ovr = inter / (vArea[i] + vArea[j] - inter);  // 计算iou
				if (ovr >= this->nms_threshold)
				{
					isSuppressed[j] = true;
				}
			}	
        }
        
    }
    int idx_t = 0;
       // remove_if()函数 remove_if(beg, end, op) //移除区间[beg,end)中每一个“令判断式:op(elem)获得true”的元素
	input_boxes.erase(remove_if(input_boxes.begin(), input_boxes.end(), [&idx_t, &isSuppressed](const BoxInfo& f) { return isSuppressed[idx_t++]; }), input_boxes.end());


}
std::vector<BoxInfo> YOLOv10::detect(Mat& frame){
    
    Mat dstimg=this->resize_image(frame);
    this->normalize_(dstimg);
    // 定义一个对应的输入举证
    array<int64_t,4> input_shape_{1,3,this->in_p_height,this->in_p_width};
    // 进行对应通道类型的转化
    cvtColor(dstimg,  dstimg, cv::COLOR_BGR2RGB);
    cv::Mat m_Normalized;
    dstimg.convertTo(m_Normalized,CV_32FC3,1/255.);
    // 进行对应的通道分离
    cv::split(m_Normalized,m_InputWrapper);
    // 创建CUDA流，推理时TensorRT是异步的，因此将CUDA流
    cudaStreamCreate(&m_CudaStream);
    auto ret = cudaMemcpyAsync(m_ArrayDevMemory[m_iInputIndex], m_ArrayHostMemory[m_iInputIndex], 
                                                                m_ArraySize[m_iInputIndex], cudaMemcpyHostToDevice, m_CudaStream); 
	// 进行异步执行，将内核排入到CUDA流中
    auto ret1 = m_CudaExecutionContext->enqueueV2(m_ArrayDevMemory, m_CudaStream, nullptr);
	ret = cudaMemcpyAsync(m_ArrayHostMemory[m_iOutputIndex], m_ArrayDevMemory[m_iOutputIndex], m_ArraySize[m_iOutputIndex], 
                            cudaMemcpyDeviceToHost, m_CudaStream); //输出传回给CPU，数据从显存到内存
    ret = cudaStreamSynchronize(m_CudaStream);
    float * pdata = (float *)m_ArrayHostMemory[m_iOutputIndex];




    std::vector<BoxInfo> generate_boxes;
    int left = {};
    int right = {};
    int top = {};
    int bottom = {};
    float confidence = {};
    int class_id = {};
    BoxInfo temp ;
    float ratiow = float(input_image_widht) / neww;
    float ratioh = float(input_image_height) / newh;
    std::cout << "conf_threshold : "<<conf_threshold << std::endl;
    for (size_t i = 0; i < 300; i++)
    {
        left = static_cast<int>(pdata[i*6 +0]);
        top = static_cast<int>(pdata[i*6 + 1]);
        right = static_cast<int>(pdata[i*6 +2]);
        bottom = static_cast<int>(pdata[i*6 + 3]);
        confidence = pdata[i*6 +4];
        class_id = static_cast<int>(pdata[i*6 +5]);
        if(confidence >= conf_threshold){
            left = static_cast<int>(left * ratiow);
            top = static_cast<int>(top * ratioh);
            right = static_cast<int>(right *ratiow);
            bottom = static_cast<int>(bottom * ratioh);
        
            temp.x1 = left;
            temp.x2 =top;
            temp.y1 = right;
            temp.y2 = bottom;
            temp.score = confidence;
            temp.label = class_id;
            temp.updatewidth();
            temp.updateLable(classes);
            generate_boxes.push_back(temp);
        }

    }
    // nms(generate_boxes);
	
    return generate_boxes;

}

void YOLOv10::init(Configuration config){
    conf_threshold = config.conf_threshold;
    nms_threshold = config.nms_threshold;
    obj_threshold = config.obj_threshold;
    classes = config.class_names;
    in_p_height = 640;
    in_p_width = 640;
    string model_path = config.model_path;
    // 加载模型
    string strTrtName = config.model_path;
    size_t sep_pos = model_path.find_last_of(".");
    strTrtName = model_path.substr(0, sep_pos) + ".engine"; // ".trt"
    bool beTure = ifFileExists(strTrtName.c_str());
	if(beTure)
    {        
        loadTrt(strTrtName);
    }else
    {
        loadOnnx(config.model_path);
    }
    // 得到输入的索引
    m_iInputIndex = m_CudaEngine->getBindingIndex("images");
    // 得到输出的索引
    m_iOutputIndex = m_CudaEngine->getBindingIndex("output0");
    // 输入
    Dims dims_i = m_CudaEngine->getBindingDimensions(m_iInputIndex);
    int size1 = dims_i.d[0]*dims_i.d[1]*dims_i.d[2]*dims_i.d[3];
    // 输入的尺度为（W H）
    m_InputSize = cv::Size(dims_i.d[3],dims_i.d[2]);
    // 输出维度(NHWC)
    Dims dims_o = m_CudaEngine->getBindingDimensions(m_iOutputIndex);
    int size2 = dims_o.d[0] * dims_o.d[1] * dims_o.d[2];   // 所有大小
    m_iClassNums = dims_o.d[2] - 5;    // [,,classes+5]
    m_iBoxNums = dims_o.d[1]; 

    // 分配内存大小
    cudaMalloc(&m_ArrayDevMemory[m_iInputIndex], size1 * sizeof(float));
    m_ArrayHostMemory[m_iInputIndex] = malloc(size1 * sizeof(float));
    m_ArraySize[m_iInputIndex] = size1 *sizeof(float);
    cudaMalloc(&m_ArrayDevMemory[m_iOutputIndex], size2 * sizeof(float));
    m_ArrayHostMemory[m_iOutputIndex] = malloc( size2 * sizeof(float));
    m_ArraySize[m_iOutputIndex] = size2 *sizeof(float);

    m_InputWrapper.emplace_back(dims_i.d[2], dims_i.d[3], CV_32FC1, m_ArrayHostMemory[m_iInputIndex]);
    m_InputWrapper.emplace_back(dims_i.d[2], dims_i.d[3], CV_32FC1, m_ArrayHostMemory[m_iInputIndex] + sizeof(float) * dims_i.d[2] * dims_i.d[3] );
    m_InputWrapper.emplace_back(dims_i.d[2], dims_i.d[3], CV_32FC1, m_ArrayHostMemory[m_iInputIndex] + 2 * sizeof(float) * dims_i.d[2] * dims_i.d[3]);
}


cv::Mat YOLOv10::YOLOv10_Draw_Line(std::vector<BoxInfo>  &generate_boxes,cv::Mat  imag){
   
    for (int i = 0; i <generate_boxes.size(); i++){

		int xmin = int(generate_boxes[i].x1);
		int ymin = int(generate_boxes[i].x2);
        int width = generate_boxes[i].y1 - generate_boxes[i].x1;
        int height = generate_boxes[i].y2 - generate_boxes[i].x2;
        cv::Rect rect(generate_boxes[i].x1,generate_boxes[i].x2,width,height);   
		rectangle(imag, rect,  Scalar(0, 0, 255), 3);
		std::string label = format("%.2f", generate_boxes[i].score);
		label = this->classes[generate_boxes[i].label] + ":" + std::to_string( generate_boxes[i].score);
		putText(imag, label, Point(xmin, ymin - 5), FONT_HERSHEY_SIMPLEX, 0.75, Scalar(0, 255, 0), 1);
    }
    return imag;
}